import os
import argparse
import torch
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from tqdm import tqdm
import re

from solar_radiation_lstm import LSTMSolarModel, load_and_preprocess_data, create_sequences

plt.rcParams['font.family'] = 'Noto Sans CJK SC'

def parse_model_dir(model_path):
    # 获取模型文件夹名
    if model_path.endswith('.pth'):
        model_dir = os.path.dirname(model_path)
    else:
        model_dir = model_path
    # 解析参数
    # 例：results_LSTM_in120_out24_h256_l2_bs512_lr0.001_drop0.2_train2019-2021_test2022
    pattern = r'in(\d+)_out(\d+)_h(\d+)_l(\d+)_bs(\d+)_lr([\d.]+)_drop([\d.]+)'
    m = re.search(pattern, model_dir)
    if not m:
        raise ValueError(f"无法从文件夹名解析模型参数: {model_dir}")
    input_window = int(m.group(1))
    output_window = int(m.group(2))
    hidden_size = int(m.group(3))
    num_layers = int(m.group(4))
    # batch_size = int(m.group(5)) # 用不到
    learning_rate = float(m.group(6)) # 用不到
    dropout = float(m.group(7))
    return {
        'input_window': input_window,
        'output_window': output_window,
        'hidden_size': hidden_size,
        'num_layers': num_layers,
        'dropout': dropout,
        'save_dir': model_dir
    }

def evaluate_and_plot(model, test_loader, scaler, device, save_dir=None, data_file=None, test_year=None, model_args=None):
    model.eval()
    predictions = []
    actuals = []
    with torch.no_grad():
        for batch_X, batch_y in tqdm(test_loader, desc="模型推理", unit="batch"):
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            outputs = model(batch_X)
            pred_reshaped = outputs.cpu().numpy().reshape(-1, 1)
            actual_reshaped = batch_y.cpu().numpy().reshape(-1, 1)
            pred_full = np.zeros((pred_reshaped.shape[0], scaler.n_features_in_))
            actual_full = np.zeros((actual_reshaped.shape[0], scaler.n_features_in_))
            pred_full[:, 0] = pred_reshaped.flatten()
            actual_full[:, 0] = actual_reshaped.flatten()
            pred_original = scaler.inverse_transform(pred_full)[:, 0]
            actual_original = scaler.inverse_transform(actual_full)[:, 0]
            predictions.extend(pred_original)
            actuals.extend(actual_original)
    predictions = np.array(predictions)
    actuals = np.array(actuals)
    # 物理约束：辐照值不能为负
    predictions[predictions < 0] = 0
    # 画图
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    axes[0].plot(actuals[:500], label='真实值', alpha=0.7)
    axes[0].plot(predictions[:500], label='预测值', alpha=0.7)
    axes[0].set_title('前500个样本预测对比')
    axes[0].legend()
    axes[0].grid(True)
    errors = predictions - actuals
    axes[1].hist(errors, bins=50, alpha=0.7, edgecolor='black')
    axes[1].set_title('预测误差分布')
    axes[1].set_xlabel('预测误差')
    axes[1].set_ylabel('频次')
    axes[1].grid(True)
    plt.tight_layout()
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        plt.savefig(os.path.join(save_dir, 'test_result.png'), dpi=300, bbox_inches='tight')
        print(f"结果图已保存到: {os.path.join(save_dir, 'test_result.png')}")
    # plt.show()
    # 打印误差
    from sklearn.metrics import mean_squared_error, mean_absolute_error
    mse = mean_squared_error(actuals, predictions)
    mae = mean_absolute_error(actuals, predictions)
    rmse = np.sqrt(mse)
    print(f"MSE: {mse:.4f}, MAE: {mae:.4f}, RMSE: {rmse:.4f}")
    # 保存csv
    if save_dir and data_file and test_year and model_args:
        df_raw = pd.read_csv(data_file)
        df_raw = df_raw[df_raw['YEAR'] == test_year].copy()
        df_raw['datetime'] = pd.to_datetime({
            'year': df_raw['YEAR'],
            'month': df_raw['MO'],
            'day': df_raw['DY'],
            'hour': df_raw['HR']
        })
        time_points = []
        months = []
        input_window = model_args['input_window']
        output_window = model_args['output_window']
        for i in range(len(actuals)):
            idx = i + input_window
            if idx < len(df_raw):
                time_points.append(df_raw.iloc[idx]['datetime'])
                months.append(df_raw.iloc[idx]['MO'])
            else:
                time_points.append(None)
                months.append(None)
        df = pd.DataFrame({
            'index': np.arange(len(actuals)),
            'datetime': time_points,
            'month': months,
            'actual': actuals,
            'pred': predictions,
            'error': predictions - actuals
        })
        csv_path = os.path.join(save_dir, 'test_result.csv')
        df.to_csv(csv_path, index=False)
        print(f"预测结果已保存到: {csv_path}")

        # 按月份画子图
        n_months = 12
        fig, axes = plt.subplots(n_months//3, 3, figsize=(18, 3*n_months//3), sharey=True)
        for m in range(1, n_months+1):
            ax = axes[(m-1)//3, (m-1)%3]
            month_df = df[df['month'] == m]
            if not month_df.empty:
                ax.plot(month_df['datetime'], month_df['actual'], label='真实值', alpha=0.7)
                ax.plot(month_df['datetime'], month_df['pred'], label='预测值', alpha=0.7)
            ax.set_title(f'{m}月')
            ax.legend()
            ax.grid(True)
            # 只显示月日时分，不显示年份
            import matplotlib.dates as mdates
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
        plt.tight_layout(rect=[0, 0, 1, 0.96])
        fig.suptitle(f'{test_year}年各月预测与真实对比', fontsize=18)
        month_fig_path = os.path.join(save_dir, 'test_result_by_month.png')
        plt.savefig(month_fig_path, dpi=300, bbox_inches='tight')
        print(f"分月对比图已保存到: {month_fig_path}")
    return predictions, actuals, mse, mae, rmse

def main():
    parser = argparse.ArgumentParser(description='LSTM太阳辐照模型测试脚本')
    parser.add_argument('--model_path', type=str, required=True, help='模型文件路径(.pth)')
    parser.add_argument('--data_file', type=str, default='POWER_Point_Hourly_20150101_20250101.csv', help='数据文件路径')
    parser.add_argument('--test_year', type=int, required=True, help='要测试的年份')
    args = parser.parse_args()

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"使用设备: {device}")

    # 自动解析模型结构参数
    model_args = parse_model_dir(args.model_path)
    print(f"自动解析到的模型参数: {model_args}")

    # 加载数据和scaler
    _, scaler, _ = load_and_preprocess_data(args.data_file, single_year=args.test_year)
    data_test_scaled, _, _ = load_and_preprocess_data(args.data_file, single_year=args.test_year, scaler=scaler)
    X_test, y_test = create_sequences(data_test_scaled, model_args['input_window'], model_args['output_window'])
    print(f"测试集序列形状: {X_test.shape}, {y_test.shape}")
    from torch.utils.data import DataLoader, Dataset
    class TempDataset(Dataset):
        def __init__(self, X, y):
            self.X = X
            self.y = y
        def __len__(self):
            return len(self.X)
        def __getitem__(self, idx):
            return torch.FloatTensor(self.X[idx]), torch.FloatTensor(self.y[idx])
    test_dataset = TempDataset(X_test, y_test)
    test_loader = DataLoader(test_dataset, batch_size=512, shuffle=False)

    # 加载模型
    input_size = X_test.shape[2]
    model = LSTMSolarModel(
        input_size=input_size,
        hidden_size=model_args['hidden_size'],
        num_layers=model_args['num_layers'],
        output_size=model_args['output_window'],
        dropout=model_args['dropout']
    ).to(device)
    model.load_state_dict(torch.load(args.model_path, map_location=device))
    print(f"模型已加载: {args.model_path}")

    # 推理与可视化
    evaluate_and_plot(
        model, test_loader, scaler, device,
        save_dir=model_args['save_dir'],
        data_file=args.data_file,
        test_year=args.test_year,
        model_args=model_args
    )

if __name__ == "__main__":
    main() 